from framework.prompt.DocLoader import DocLoader

class PromptController():
    def __init__(self, DocLoader:DocLoader, template:str) -> None:
        self.template = template
        self.DocLoader = DocLoader
        if self.DocLoader.vectostore:
            self.retriever = self.DocLoader.vectostore.as_retriever()
        else:
            self.retriever = None
    
    def gen_message(self,query:str):
        if self.retriever:
            context = self.retriever.invoke(query)
            print(f"context=$$\n\n{self.format_docs(context)}\n\n$$")
        else:
            raise ValueError("retriever not found, please check whether you have added documents by DocLoader to create an vectorstore")
        context_str = self.format_docs(context)
        prompt = self.template.format(context=context_str, question=query)
        message = [{"role":"user", "content":prompt}]
        return message
    
    def format_docs(self,docs):
        return "\n\n".join(doc.page_content for doc in docs)
    